import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import shap
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import cross_val_score


# Ensure the picture folder exists
def ensure_picture_folder_exists():
    """Ensure the picture folder exists."""
    picture_folder = "./picture"
    if not os.path.exists(picture_folder):
        os.makedirs(picture_folder)
    return picture_folder

# 1. Load Data
def load_data(filepath):
    """Load data from a CSV file."""
    try:
        data = pd.read_csv(filepath, encoding='ISO-8859-1')
        print(f"Data file '{filepath}' read successfully!")
        return data
    except FileNotFoundError:
        print(f"Error: The file '{filepath}' was not found. Please check the file path!")
        exit()
    except UnicodeDecodeError:
        print(f"Error: The file '{filepath}' encoding is not supported by ISO-8859-1. Please try another encoding format!")
        exit()

# 2. Data Preprocessing
def preprocess_data(data):
    """Preprocess the data by filling missing values and converting categorical features."""
    # Fill missing values with 0
    data.fillna(0, inplace=True)
    # Convert country code (NOC) to numerical feature
    data['NOC'] = data['NOC'].astype('category').cat.codes
    return data

# 3. Feature Extraction
def extract_features(data):
    """Extract features and target variables from the data."""
    features = ['Year', 'NOC', 'Gold', 'Silver', 'Bronze', 'Total']
    X = data[features]
    y_gold = data['Gold']  # Target: Predict gold medals
    y_total = data['Total']  # Target: Predict total medals
    return X, y_gold, y_total

# 4. Data Visualization
def visualize_data(data, picture_folder):
    """Visualize the data using various plots."""
    # 4.1 Distribution of medal counts
    plt.figure(figsize=(10, 6))
    sns.histplot(data['Gold'], bins=30, kde=True, color='blue', label='Gold Medals')
    sns.histplot(data['Total'], bins=30, kde=True, color='red', label='Total Medals')
    plt.title('Distribution of Gold and Total Medals')
    plt.xlabel('Medal Count')
    plt.ylabel('Frequency')
    plt.legend()
    plt.grid(axis='y')  # Add vertical grid lines only
    plt.savefig(f"{picture_folder}/1.2.1.png")
    plt.show()

    # 4.2 Relationship between gold medals and total medals
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x=data['Gold'], y=data['Total'], color='purple', alpha=0.6)
    sns.regplot(x=data['Gold'], y=data['Total'], color='orange', scatter=False)
    plt.title('Relationship between Gold and Total Medals')
    plt.xlabel('Gold Medals')
    plt.ylabel('Total Medals')
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.2.png")
    plt.show()

    # 4.3 Medal trends over years
    plt.figure(figsize=(10, 6))
    sns.lineplot(x=data['Year'], y=data['Gold'], color='blue', label='Gold Medals')
    sns.lineplot(x=data['Year'], y=data['Total'], color='red', label='Total Medals')
    plt.title('Medal Trends Over Years')
    plt.xlabel('Year')
    plt.ylabel('Medal Count')
    plt.legend()
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.3.png")
    plt.show()

# 5. Model Training
def train_models(X, y_gold, y_total):
    """Train Random Forest Regressor models and evaluate their performance."""
    model_gold = RandomForestRegressor(n_estimators=100, random_state=42)
    model_total = RandomForestRegressor(n_estimators=100, random_state=42)

    # Cross-validation to evaluate model performance
    cv_scores_gold = cross_val_score(model_gold, X, y_gold, cv=5, scoring='neg_mean_squared_error')
    cv_scores_total = cross_val_score(model_total, X, y_total, cv=5, scoring='neg_mean_squared_error')

    print(f"Cross-validation MSE for gold medal prediction: {-cv_scores_gold.mean():.2f}")
    print(f"Cross-validation MSE for total medal prediction: {-cv_scores_total.mean():.2f}")

    # Train models
    model_gold.fit(X, y_gold)
    model_total.fit(X, y_total)

    return model_gold, model_total

# 6. Model Visualization
def visualize_model_performance(model_gold, model_total, X, y_gold, y_total, picture_folder):
    """Visualize the model performance and feature importance."""
    # 6.1 Feature importance for gold medal prediction
    plt.figure(figsize=(10, 6))
    plt.barh(X.columns, model_gold.feature_importances_, color='skyblue')
    plt.title('Feature Importance for Gold Medal Prediction')
    plt.xlabel('Importance')
    plt.ylabel('Feature')
    plt.grid(axis='x')  # Add vertical grid lines only
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.2.4.png")
    plt.show()

    # 6.2 Predictions
    y_pred_gold = model_gold.predict(X)  # Predicted gold medals
    y_pred_total = model_total.predict(X)  # Predicted total medals

    # 6.3 Actual vs Predicted Gold Medals
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x=y_gold, y=y_pred_gold, color='blue', alpha=0.6)
    plt.plot([min(y_gold), max(y_gold)], [min(y_gold), max(y_gold)], color='red', linestyle='--')
    plt.title('Actual vs Predicted Gold Medals')
    plt.xlabel('Actual Gold Medals')
    plt.ylabel('Predicted Gold Medals')
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.5.png")
    plt.show()

    # 6.4 Actual vs Predicted Total Medals
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x=y_total, y=y_pred_total, color='green', alpha=0.6)
    plt.plot([min(y_total), max(y_total)], [min(y_total), max(y_total)], color='red', linestyle='--')
    plt.title('Actual vs Predicted Total Medals')
    plt.xlabel('Actual Total Medals')
    plt.ylabel('Predicted Total Medals')
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.6.png")
    plt.show()

# 7. Predict Medals for 2028
def predict_medals_for_2028(model_gold, model_total, data):
    """Predict the number of medals for the 2028 Olympics."""
    countries = data['NOC'].unique()
    predictions = []

    for country in countries:
        # Construct features for 2028
        X_2028 = pd.DataFrame({
            'Year': [2028],
            'NOC': [country],
            'Gold': [data[data['NOC'] == country]['Gold'].mean()],  # Use historical average
            'Silver': [data[data['NOC'] == country]['Silver'].mean()],
            'Bronze': [data[data['NOC'] == country]['Bronze'].mean()],
            'Total': [data[data['NOC'] == country]['Total'].mean()]
        })

        # Predict gold medals and total medals
        gold_pred = model_gold.predict(X_2028)[0]
        total_pred = model_total.predict(X_2028)[0]

        predictions.append({
            'NOC': country,
            'Gold_2028': gold_pred,
            'Total_2028': total_pred
        })

    predictions_df = pd.DataFrame(predictions)
    print("Predictions for 2028 Olympics:")
    print(predictions_df)
    return predictions_df

# 8. SHAP Analysis
def shap_analysis(model_gold, model_total, X, picture_folder):
    """Perform SHAP analysis for model explainability."""
    # Initialize SHAP explainer
    explainer_gold = shap.TreeExplainer(model_gold)
    shap_values_gold = explainer_gold.shap_values(X)

    explainer_total = shap.TreeExplainer(model_total)
    shap_values_total = explainer_total.shap_values(X)

    # 8.1 SHAP summary plot for gold medal prediction
    plt.figure(figsize=(10, 6))
    shap.summary_plot(shap_values_gold, X, feature_names=X.columns, show=False)
    plt.title('SHAP Summary Plot for Gold Medal Prediction')
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.7.png")
    plt.show()

    # 8.2 SHAP summary plot for total medal prediction
    plt.figure(figsize=(10, 6))
    shap.summary_plot(shap_values_total, X, feature_names=X.columns, show=False)
    plt.title('SHAP Summary Plot for Total Medal Prediction')
    plt.grid(True)  # Add both horizontal and vertical grid lines
    plt.savefig(f"{picture_folder}/1.2.8.png")
    plt.show()

    # 8.3 SHAP force plot for a single sample (gold medal prediction)
    sample_idx = 0  # Select the first sample
    shap.force_plot(explainer_gold.expected_value, shap_values_gold[sample_idx, :], X.iloc[sample_idx, :], feature_names=X.columns)

# Main execution
if __name__ == "__main__":
    # Ensure the picture folder exists
    picture_folder = ensure_picture_folder_exists()

    # Step 1: Load data
    data = load_data(r'summerOly_medal_counts.csv')

    # Step 2: Data preprocessing
    data = preprocess_data(data)

    # Step 3: Feature extraction
    X, y_gold, y_total = extract_features(data)

    # Step 4: Data visualization
    visualize_data(data, picture_folder)

    # Step 5: Model training
    model_gold, model_total = train_models(X, y_gold, y_total)

    # Step 6: Model visualization
    visualize_model_performance(model_gold, model_total, X, y_gold, y_total, picture_folder)

    # Step 7: Predict medals for 2028
    predictions_2028 = predict_medals_for_2028(model_gold, model_total, data)

    # Step 8: SHAP analysis
    shap_analysis(model_gold, model_total, X, picture_folder)